from langgraph.graph import StateGraph,START,END
from langchain_core.messages import HumanMessage,AIMessage,SystemMessage
from langchain_openai import ChatOpenAI
import os
# from langchain_core.prompts import ChatPromptTemplate
from typing import Optional,Union
from pydantic import BaseModel,Field
llm = ChatOpenAI(
  base_url='http://localhost:11434/v1',
  api_key=os.environ['DEEPSEEK_API_KEY'],
  model = 'qwen3:8b',
)
# prompt=ChatPromptTemplate.from_messages(
#   [('system','Answer the user query. Wrap the output in `json`'),
#   ('human','{input}'),]
# )
class UersInfo(BaseModel):
  '''
  负责对用户信息的提取，包括姓名、年龄、地址、手机号和邮箱地址。
  '''
  name:str=Field(description='用户姓名') # 必填
  age:Optional[int]=Field(description='用户年龄') # 可选
  address:Optional[str]=Field(description='用户地址') # 可选
  phone:str=Field(description='用户手机号') # 必填
  email:Optional[str]=Field(description='用户邮箱地址') # 可选
class ConversationResponse(BaseModel):
  '''
  定义一个模型，用于处理用户输入的文本。
  '''
  response:str=Field(description='用于对用户输入进行自然语言的回答')
class FinalModel(BaseModel):
  '''
  定义一个模型，用于选择模型运行规则。
  '''
  final_output:Union[UersInfo,ConversationResponse]=Field(description='根据用户输入选择的模型运行规则')

structured_llm = llm.with_structured_output(FinalModel)

res = structured_llm.invoke('我叫zyy，我今年18岁，我来自中国，邮箱是zyy@example.com，手机号是13800000000')

print(res)

